import os
import os.path as osp
from tqdm import tqdm
import json
from PIL import Image
import numpy as np
import cv2
import shutil
if __name__ == '__main__':

    import random
    jsondir = '/llmcapagroup1/test-bucket/liuxin/shiyout2i/prompt/328/yingxiao_v1_v4_shiyoudata15w_doubao18w.json'
    datalist = []
    data = open(jsondir, 'r', encoding='utf-8')
    for line in tqdm(data.readlines()):
        dic = json.loads(line)
        datalist.append(dic)
    print('datalist len:', len(datalist))

    targetdictv1 = {'768x1280':0.6, '960x1280':0.75, '1024x1024':1, '560x1280':0.4375, '704x1280':0.55, '848x1280':0.6625}  # hxw:h/w, 横图
    targetdictv2 = {'1280x960':1.3333, '1280x768':1.6667, '1280x832':1.5385, '1280x800':1.6, '1280x576':2.2222} # 竖图
    newdict = {'768x1280':[], '960x1280':[], '1024x1024':[],'1280x960':[], '1280x768':[], '560x1280':[], '704x1280':[], '848x1280':[], '1280x832':[], '1280x800':[], '1280x576':[]}
    for i, f in enumerate(datalist):
        h, w = f['size'] # h,w
        h, w = int(h), int(w)
        ratio = int(h) / int(w)
        if h <= w:
            res_key, res_val = min(targetdictv1.items(), key=lambda x: abs(ratio - x[1]))
        else:
            res_key, res_val = min(targetdictv2.items(), key=lambda x: abs(ratio - x[1]))
        f['target_size'] = res_key
        newdict[res_key].append(f)
    
    savelist = []
    for k,v in newdict.items():
        v = v[0:(len(v)//32)*32] # 单卡batchsize*单卡gpus*机器数  32=4*8*1 
        random.shuffle(v)
        print(k, len(v))
        savelist += v[0:45]
    print(len(savelist))


    # datasamplelist = random.sample(datalist, 500)
    datasamplelist = savelist

    savedir = 'data/imgs'
    os.makedirs(savedir, exist_ok=True)
    savejsonfile = 'data/train_data.json'

    with open(savejsonfile, 'w', encoding='utf-8') as f:
        for index, item in enumerate(tqdm(datasamplelist)):
            imgpath = item['imgpath']
            newpath = osp.join(savedir, str(index) + '.jpg')
            shutil.copy(imgpath, newpath)
            dic = {'imgpath':newpath, 'text_zh':"艺术海报, 高质量", 'size':item['size']}
            if osp.exists(dic['imgpath']):
                f.write(json.dumps(dic, ensure_ascii=False) + '\n')